package crypto

import (
	"crypto/rand"
	"fmt"
	"math/big"
)

// PaillierPublicKey represents a Paillier public key
type PaillierPublicKey struct {
	N       *big.Int // n = p * q
	NSquare *big.Int // n^2
	G       *big.Int // generator, typically n + 1
}

// PaillierPrivateKey represents a Paillier private key
type PaillierPrivateKey struct {
	Lambda *big.Int // λ = lcm(p-1, q-1)
	Mu     *big.Int // μ = L(g^λ mod n^2)^-1 mod n
	Public *PaillierPublicKey
}

// PaillierCiphertext represents an encrypted value
type PaillierCiphertext struct {
	C *big.Int // ciphertext
}

// GeneratePaillierKeyPair generates a new Paillier key pair
func GeneratePaillierKeyPair(bits int) (*PaillierPublicKey, *PaillierPrivateKey, error) {
	// Generate two large primes p and q
	p, err := rand.Prime(rand.Reader, bits/2)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate prime p: %w", err)
	}

	q, err := rand.Prime(rand.Reader, bits/2)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate prime q: %w", err)
	}

	// Ensure p != q
	for p.Cmp(q) == 0 {
		q, err = rand.Prime(rand.Reader, bits/2)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to generate distinct prime q: %w", err)
		}
	}

	// Compute n = p * q
	n := new(big.Int).Mul(p, q)
	nSquare := new(big.Int).Mul(n, n)

	// Use g = n + 1 as generator (simplified version)
	g := new(big.Int).Add(n, big.NewInt(1))

	publicKey := &PaillierPublicKey{
		N:       n,
		NSquare: nSquare,
		G:       g,
	}

	// Compute λ = lcm(p-1, q-1)
	pMinus1 := new(big.Int).Sub(p, big.NewInt(1))
	qMinus1 := new(big.Int).Sub(q, big.NewInt(1))
	lambda := new(big.Int).Mul(pMinus1, qMinus1)
	lambda.Div(lambda, new(big.Int).GCD(nil, nil, pMinus1, qMinus1))

	// Compute μ = L(g^λ mod n^2)^-1 mod n
	gLambda := new(big.Int).Exp(g, lambda, nSquare)
	mu := lFunction(gLambda, n)
	mu.ModInverse(mu, n)

	privateKey := &PaillierPrivateKey{
		Lambda: lambda,
		Mu:     mu,
		Public: publicKey,
	}

	return publicKey, privateKey, nil
}

// lFunction computes L(x) = (x-1)/n
func lFunction(x, n *big.Int) *big.Int {
	result := new(big.Int).Sub(x, big.NewInt(1))
	result.Div(result, n)
	return result
}

// Encrypt encrypts a plaintext using Paillier encryption
func (pk *PaillierPublicKey) Encrypt(plaintext *big.Int, r *big.Int) (*PaillierCiphertext, error) {
	if r == nil {
		// Generate random r
		var err error
		r, err = rand.Int(rand.Reader, pk.N)
		if err != nil {
			return nil, fmt.Errorf("failed to generate random r: %w", err)
		}
	}

	// Ensure r is coprime to n
	for new(big.Int).GCD(nil, nil, r, pk.N).Cmp(big.NewInt(1)) != 0 {
		var err error
		r, err = rand.Int(rand.Reader, pk.N)
		if err != nil {
			return nil, fmt.Errorf("failed to generate coprime r: %w", err)
		}
	}

	// c = g^m * r^n mod n^2
	gM := new(big.Int).Exp(pk.G, plaintext, pk.NSquare)
	rN := new(big.Int).Exp(r, pk.N, pk.NSquare)
	c := new(big.Int).Mul(gM, rN)
	c.Mod(c, pk.NSquare)

	return &PaillierCiphertext{C: c}, nil
}

// Decrypt decrypts a ciphertext using Paillier decryption
func (sk *PaillierPrivateKey) Decrypt(ciphertext *PaillierCiphertext) (*big.Int, error) {
	// m = L(c^λ mod n^2) * μ mod n
	cLambda := new(big.Int).Exp(ciphertext.C, sk.Lambda, sk.Public.NSquare)
	m := lFunction(cLambda, sk.Public.N)
	m.Mul(m, sk.Mu)
	m.Mod(m, sk.Public.N)

	return m, nil
}

// Add performs homomorphic addition: E(a) + E(b) = E(a + b)
func (pk *PaillierPublicKey) Add(c1, c2 *PaillierCiphertext) *PaillierCiphertext {
	result := new(big.Int).Mul(c1.C, c2.C)
	result.Mod(result, pk.NSquare)
	return &PaillierCiphertext{C: result}
}

// Mul performs homomorphic multiplication by a constant: E(a) * k = E(a * k)
func (pk *PaillierPublicKey) Mul(ciphertext *PaillierCiphertext, k *big.Int) *PaillierCiphertext {
	result := new(big.Int).Exp(ciphertext.C, k, pk.NSquare)
	return &PaillierCiphertext{C: result}
}

// GetRandomLtN returns a random number less than n
func (pk *PaillierPublicKey) GetRandomLtN() *big.Int {
	r, err := rand.Int(rand.Reader, pk.N)
	if err != nil {
		panic(fmt.Sprintf("failed to generate random number: %v", err))
	}
	return r
}

// Ciphertext returns the ciphertext value
func (ct *PaillierCiphertext) Ciphertext(be_secure bool) *big.Int {
	// For simplicity, we ignore the be_secure parameter
	// In a full implementation, this would add randomization if be_secure is true
	return new(big.Int).Set(ct.C)
}

// RawAdd performs raw addition of ciphertexts (for homomorphic operations)
func (ct *PaillierCiphertext) RawAdd(other *PaillierCiphertext) *big.Int {
	result := new(big.Int).Add(ct.C, other.C)
	return result
}
